import asyncio
import os
import sys
import threading
from pathlib import Path
import cv2
import serial
import serial.tools.list_ports
import torch
import torch.backends.cudnn as cudnn
import yaml
from PyQt5 import QtWidgets
from PyQt5.QtCore import *
from PyQt5.QtGui import *
from PyQt5.QtMultimedia import QMediaPlayer, QMediaContent
from PyQt5.QtWidgets import *
from gui import Ui_Form
from models.common import DetectMultiBackend
from utils.datasets import LoadImages, LoadStreams
from utils.general import (LOGGER, check_img_size, non_max_suppression, scale_coords, xyxy2xywh)
from utils.plots import Annotator, colors
from utils.torch_utils import select_device, time_sync

FILE = Path(__file__).resolve()
# 这里是有关根目录的
ROOT = FILE.parents[0]  # YOLOv5 root directory
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative


# python -m PyQt5.uic.pyuic gui.ui -o gui.py  转换代码

class MainWindows(QtWidgets.QWidget, Ui_Form):
    def __init__(self):
        super(MainWindows, self).__init__()
        # 这里是读取 每种类型的垃圾对应的种类
        self.yhlj = ['battery', 'drug']
        self.khslj = ['can', 'bottle']
        self.cylj = ['rabbit', 'potato', 'mooli', 'cucumber']
        self.qtlj = ['china', 'brick', 'cobblestone', 'eggplat']
        self.names = ['yhlj', 'khslj', 'qtlj', 'cylj']
        self.name_dict = {}
        for name in self.names:
            self.name_dict[name] = 0
        self.setupUi(self)
        self.ser = serial.Serial()  # 打开串口按钮
        self.port_open()  # 定时器接收数据
        self.timer = QTimer(self)
        self.pushButton_streaming.clicked.connect(self.open_cam)
        self.pushButton_stopscan.clicked.connect(self.stop_cam)
        self.setWindowTitle("垃圾分类")
        self.setWindowIcon(QIcon("images/UI/logo1.jpg"))
        self.output_size = 1080  # 图片读取进程
        self.device = 'cuda:0'
        self.vid_source = 0  # 初始设置为摄像头
        # self.stopEvent = threading.Event()
        # self.stopEvent.clear()
        self.webcam = False
        self.in_detect = False  # 是否正在检测
        self.model = self.model_load(weights="pretrained/best.pt",
                                     device=self.device, dnn=False)  # todo 指明模型加载的位置的设备
        self.reset()
        self.player = QMediaPlayer()  # 播放器
        self.player.setVideoOutput(self.wgt_player)
        self.player.mediaStatusChanged.connect(self.handle_media_status_changed)
        self.video_path = os.path.abspath('images/ljfl.mp4')
        self.player.setMedia(QMediaContent(QUrl.fromLocalFile(self.video_path)))
        self.player.play()
        self.is_send = False
        self.no_send = 0

    def handle_media_status_changed(self, status):
        if status == self.player.EndOfMedia:
            # 当视频播放完成时，重新播放视频
            self.player.setPosition(0)
            self.player.play()

    @torch.no_grad()  # 一个装饰器，用于修饰下面的函数或方法，表示在该函数或方法内部不进行梯度计算。这通常用于 PyTorch 中的推断或评估阶段，以减少内存使用和提高性能。
    def model_load(self, weights="",  # model.pt path(s)
                   device='',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
                   half=False,  # use FP16 half-precision inference
                   dnn=False,  # use OpenCV DNN for ONNX inference
                   ):
        '''模型初始化'''
        # 这段代码的目的是选择并配置深度学习模型的设备（device），通常用于 PyTorch 或类似深度学习框架的代码中。
        device = select_device(device)
        half &= device.type != 'cpu'  # half precision only supported on CUDA
        device = select_device(device)
        model = DetectMultiBackend(weights, device=device, dnn=dnn)
        stride, names, pt, jit, onnx = model.stride, model.names, model.pt, model.jit, model.onnx
        # Half
        half &= pt and device.type != 'cpu'  # half precision only supported by PyTorch on CUDA
        if pt:
            model.model.half() if half else model.model.float()
        print("模型加载完成!")
        return model

    def reset(self):
        """这个方法执行了一些操作，用于重置类的某些状态或界面元素。"""
        self.pushButton_streaming.setEnabled(True)
        self.pushButton_stopscan.setEnabled(False)
        self.label_video.setPixmap(QPixmap("images/UI/yolo.jpg"))
        self.vid_source = '0'
        self.webcam = True

    def closeEvent(self, event):
        """界面关闭事件"""
        reply = QMessageBox.question(self,
                                     '退出',
                                     "Are you sure?",
                                     QMessageBox.Yes | QMessageBox.No,
                                     QMessageBox.No)
        if reply == QMessageBox.Yes:
            # 关闭摄像头
            self.stop_cam()
            self.close()
            event.accept()
            print('系统退出')
            sys.exit()
        else:
            event.ignore()

    def open_cam(self):
        """点击开启摄像头就跳转到这里"""
        self.pushButton_streaming.setEnabled(False)  # 打开按钮设置为不可按
        self.pushButton_stopscan.setEnabled(True)
        self.vid_source = '0'
        self.webcam = True
        if self.in_detect == False:
            self.in_detect = True  # 必须放在asyncio前面
            print('开始检测')
            asyncio.get_event_loop().run_until_complete(self.detect())
            # asyncio.run(self.detect()) # 适用于python3.7以后的版本
        else:
            print('已经在检测了')
        # 启动摄像头线程
        # self.stopEvent.clear()
        # self.camera_thread = threading.Thread(target=self.detect) # 使用threading多线程的写法
        # self.camera_thread.start()

    def stop_cam(self):
        """关闭摄像头"""
        self.pushButton_streaming.setEnabled(True)
        self.pushButton_stopscan.setEnabled(False)
        # 停止摄像头线程
        if self.in_detect == True:
            self.in_detect = False
            print('停止检测')
        else:
            print('没有检测，无法停止')
        # self.stopEvent.set() # 使用threading多线程的写法
        # self.camera_thread.join() # 使用threading多线程的写法

    async def detect(self):
        """检测代码"""
        model = self.model
        output_size = self.output_size  # 输出大小
        imgsz = [320, 320]  # inference size (pixels) 电脑是1280*720 640,640会出问题，到1280不会了
        conf_thres = 0.6  # confidence threshold
        iou_thres = 0.45  # NMS IOU threshold
        max_det = 1  # maximum detections per image
        device = self.device  # cuda device, i.e. 0 or 0,1,2,3 or cpu
        view_img = True  # show results，这个很重要，原来输出给GUI
        save_txt = False  # save results to *.txt
        save_conf = False  # save confidences in --save-txt labels
        save_crop = False  # save cropped prediction boxes
        nosave = False  # do not save images/videos
        classes = None  # filter by class: --class 0, or --class 0 2 3
        agnostic_nms = False  # class-agnostic NMS
        augment = False  # ugmented inference
        visualize = False  # visualize features
        # 这些在代码在前面加载模型部分
        # update = False,  # update all models
        # project = ROOT / 'runs/detect',  # save results to project/name
        # name = 'exp',  # save results to project/name
        # exist_ok = False,  # existing project/name ok, do not increment
        line_thickness = 3  # bounding box thickness (pixels)
        hide_labels = False  # hide labels
        hide_conf = False  # hide confidences
        half = False  # use FP16 half-precision inference
        dnn = False  # use OpenCV DNN for ONNX inference
        # TODO:上面是yolo自己的变量，下面是我添加的
        source = str(self.vid_source)
        webcam = self.webcam

        # save_img = not nosave and not source.endswith('.txt')  # save inference images
        # is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)
        # is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://'))
        # webcam = source.isnumeric() or source.endswith('.txt') or (is_url and not is_file)
        # if is_url and is_file:
        #     source = check_file(source)  # download
        #
        # # Directories
        # save_dir = increment_path(Path(project) / name, exist_ok=exist_ok)  # increment run
        # (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir

        # Load model
        device = select_device(self.device)
        stride, names, pt, jit, onnx = model.stride, model.names, model.pt, model.jit, model.onnx
        imgsz = check_img_size(imgsz, s=stride)  # check image size
        save_img = not nosave and not source.endswith('.txt')  # save inference images

        # Dataloader
        if webcam:
            # view_img = check_imshow() # 检测是否可以使用imshow展示
            cudnn.benchmark = True  # set True to speed up constant image size inference
            dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt and not jit)
            bs = len(dataset)  # batch_size
        else:
            dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt and not jit)
            bs = 1  # batch_size
        vid_path, vid_writer = [None] * bs, [None] * bs
        # Run inference
        if pt and device.type != 'cpu':
            # 明在*imgsz中的imgsz不是一个可迭代对象，而是一个整数。所以imgsz=(imgsz,),然后GPT再指出，少了一个维度。
            model(torch.zeros(1, 3, *imgsz).to(device).type_as(next(model.model.parameters())))  # warmup
        dt, seen = [0.0, 0.0, 0.0], 0
        for path, im, im0s, vid_cap, s in dataset:
            is_ok = True # 这次检测是否合理，如果不合理，后面会不保存本次结果
            t1 = time_sync()

            # _, _, height, width = im.shape
            # crop_x1 = 3 * width // 8  # 左上角 x 坐标
            # crop_x2 = 5 * width // 8  # 右下角 x 坐标
            # crop_y1 = 3 * height // 8  # 左上角 y 坐标
            # crop_y2 = 5 * height // 8  # 右下角 y 坐标
            # im = im[:, :, crop_y1:crop_y2, crop_x1:crop_x2]
            # cropped_im0s = []
            # for img in im0s:
            #     cropped_img = img[crop_y1:crop_y2, crop_x1:crop_x2]
            #     cropped_im0s.append(cropped_img)

            im = torch.from_numpy(im).to(device)  # 发送到device
            im = im.half() if half else im.float()  # uint8 to fp16/32
            im /= 255  # 0 - 255 to 0.0 - 1.0 # 把默认8位的数字归到0-1之间
            if len(im.shape) == 3:
                im = im[None]  # expand for batch dim
            t2 = time_sync()
            dt[0] += t2 - t1

            # Inference
            # visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False
            pred = model(im, augment=augment, visualize=visualize)
            t3 = time_sync()
            dt[1] += t3 - t2

            # NMS
            pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
            dt[2] += time_sync() - t3

            # Second-stage classifier (optional)
            # pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)

            # Process predictions
            for i, det in enumerate(pred):  # per image
                seen += 1
                if webcam:  # batch_size >= 1
                    p, im0, frame = path[i], im0s[i].copy(), dataset.count
                    s += f'{i}: '
                else:
                    p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)
                p = Path(p)  # to Path
                # save_path = str(save_dir / p.name)  # im.jpg
                # txt_path = str(save_dir / 'labels' / p.stem) + (
                #     '' if dataset.mode == 'image' else f'_{frame}')  # im.txt
                s += '%gx%g ' % im.shape[2:]  # print string
                gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
                imc = im0.copy() if save_crop else im0  # for save_crop
                cv2.imwrite("images/tmp/single_org_vid.jpg", im0)
                annotator = Annotator(im0, line_width=line_thickness, example=str(names))
                if len(det):
                    # Rescale boxes from img_size to im0 size
                    det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()

                    # Print results
                    for c in det[:, -1].unique():
                        n = (det[:, -1] == c).sum()  # detections per class
                        s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string

                    print('本次检测到物体数量：', len(det))
                    print('本次检测到物体的坐标：', det[0][4])
                    # 实现在所有标注框中只标注置信度最大的det_max
                    MaxConf = []
                    len_det = len(det)
                    for i in range(len_det):
                        # print(i)
                        MaxConf.append(det[i][4])
                        Max = max(MaxConf)
                        # 最大值
                        MaxI = MaxConf.index(Max)
                    det_max = det[[MaxI]]
                    # Write results
                    for *xyxy, conf, cls in reversed(det):  # 每一个目标
                        if save_txt:  # Write to file
                            xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(
                                -1).tolist()  # normalized xywh
                            line = (cls, *xywh, conf) if save_conf else (cls, *xywh)  # label format
                            # with open(txt_path + '.txt', 'a') as f:
                            #     f.write(('%g ' * len(line)).rstrip() % line + '\n')

                        if save_img or save_crop or view_img:  # Add bbox to image
                            c = int(cls)  # integer class
                            label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
                            # annotator.box_label(xyxy, label, color=colors(c, True))
                            # if save_crop:
                            #     save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)

                            # TODO:发送字符串代码和显示图片，以上部分是yolo自带，下面是我们写的
                            str_box = annotator.box_label(xyxy, label, color=colors(c, True))
                            self.textBrowser_video.setText(str_box)
                            t = str(xyxy) + str(label)
                            # with open("ARM_Controller/temp.txt", "w") as f:
                            #     f.write(t + "\n")

                            # if self.continue_detection:
                            #     self.continue_detection = False
                            #     await asyncio.ensure_future(self.save_send(label, xyxy, gn))
                            # else:
                            #     print('等待反馈，跳过信息发送')
                            try:
                                xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(
                                    -1).tolist()  # normalized xywh
                                # 坐标信息（假设 x=100，y=100，角度为0°,设置摄像头视野大小为1000）
                                x_center = int(xywh[0] * 256)
                                y_center = int(xywh[1] * 256)
                                x_left = int((xywh[0] - xywh[2] / 2) * 256)
                                x_right = int((xywh[0] + xywh[2] / 2) * 256)
                                y_up = int((xywh[1] - xywh[3] / 2) * 256)
                                y_down = int((xywh[1] + xywh[3] / 2) * 256)
                                w = int(xywh[2] * 256)
                                h = int(xywh[3] * 256)
                                if x_center < 32 or x_center > 224 or y_center < 32 or y_center > 224 or w > 128 or h > 128:
                                    print('-'*10 + '这次检测结果在边缘或者过大，跳过结果' + '-'*10)
                                    is_ok = False
                                    continue  # 中心点在外面

                                if self.is_send or self.no_send >= 100:
                                    if self.no_send >= 100:
                                        print('\n' * 5 + '-' * 20 + '100次没有收到反馈，再次发送信号' + '-' * 20)
                                    garbage_type = 1
                                    label = label.split()[0]  # 分离label，物体的类型
                                    # 给物体计数，在UI界面显示
                                    if label in self.yhlj:
                                        garbage_type = 1
                                        self.name_dict['yhlj'] += 1
                                    if label in self.khslj:
                                        garbage_type = 2
                                        self.name_dict['khslj'] += 1
                                    if label in self.cylj:
                                        garbage_type = 3
                                        self.name_dict['cylj'] += 1
                                    if label in self.qtlj:
                                        garbage_type = 4
                                        self.name_dict['qtlj'] += 1
                                    self.label_cy.setText(str(self.name_dict['cylj']))
                                    self.label_yh.setText(str(self.name_dict['yhlj']))
                                    self.label_khs.setText(str(self.name_dict['khslj']))
                                    self.label_qt.setText(str(self.name_dict['qtlj']))
                                    # 这一块是检测是否满载
                                    self.label_cym.setText(str('未满'))
                                    self.label_yhm.setText(str('未满'))
                                    self.label_khsm.setText(str('未满'))
                                    self.label_qtm.setText(str('未满'))
                                    if (xywh[0] - xywh[1]) < 0.1:
                                        angle = 0
                                    elif (xywh[2] - xywh[3]) < 0.1:
                                        angle = 1
                                    else:
                                        angle = 2
                                    angle = 1  # 0°xywh[2] + xywh[3]
                                    # 将垃圾种类和坐标信息编码为二进制
                                    binary_data = bytearray([
                                        0xFF, 0xFE,  # 开头标志
                                        0x01,  # 单垃圾分类
                                        garbage_type,  # 垃圾种类
                                        x_center,
                                        y_center,
                                        # x_center >> 8, x_center & 0xFF,  # x轴坐标（高位在前，低位在后）
                                        # y_center >> 8, y_center & 0xFF,  # y轴坐标（高位在前，低位在后）
                                        angle,  # 角度
                                        0xFC  # 结尾标志
                                    ])
                                    self.ser.write(binary_data)  # 将字符串编码为字节并发送
                                    # 下次是否需要发送置False，把没有发送信息次数置零
                                    self.is_send = False
                                    self.no_send = 0
                                else:
                                    self.no_send += 1
                                # host_data = ' '.join([format(byte, '02X') for byte in binary_data]) # 将二进制数据转换为十六进制字符串
                                # self.ser.write(host_data.encode())  # 将字符串编码为字节并发送
                                # 等待对方返回
                                self.ser.timeout = 0.01  # 设置读取超时时间为0.01秒
                                response = self.ser.read()  # .decode()
                                response = response.hex()
                                # 根据返回值进行记录
                                if response == '01':
                                    print("串口反馈01：分类成功")
                                    self.is_send = True
                                    # with open("records.txt", "a") as file:  # 保存记录到文本文件
                                    #     file.write(xywh + "\n")
                                elif response == '00':
                                    self.is_send = True
                                    print("串口反馈00：分类失败", response)
                                    # with open("records.txt", "a") as file:  # 在这里执行记录失败的操作
                                    #     file.write("抓取失败\n")
                                else:
                                    print("未知的响应或者无反馈")
                            except serial.SerialException as e:
                                print(f"串口通信错误: {e}")  #

                # Print time (inference-only)
                LOGGER.info(f'{s}Done. ({t3 - t2:.3f}s)')

                # Stream results（yolov5结果显示）
                im0 = annotator.result()
                frame = im0
                resize_scale = output_size / frame.shape[0]
                frame_resized = cv2.resize(frame, (0, 0), fx=resize_scale, fy=resize_scale)
                if is_ok:
                    cv2.imwrite("images/tmp/single_result_vid.jpg", frame_resized)
                    self.label_video.setPixmap(QPixmap("images/tmp/single_result_vid.jpg"))
                else:
                    self.label_video.setPixmap(QPixmap("images/tmp/single_org_vid.jpg"))
                self.label_video.setScaledContents(True)
                # if view_img: #这是yolo原来的结果显示，使用的cv2
                #     cv2.imshow(str(p), im0)
                #     cv2.waitKey(1)  # 1 millisecond

                # Save results (image with detections)  #这里是detect的保存图片部分，但是目前我们使用GUI不保存结果，保存图片的路径也删除了
                # if save_img:
                #     if dataset.mode == 'image':
                #         cv2.imwrite(save_path, im0)
                #     else:  # 'video' or 'stream'
                #         if vid_path[i] != save_path:  # new video
                #             vid_path[i] = save_path
                #             if isinstance(vid_writer[i], cv2.VideoWriter):
                #                 vid_writer[i].release()  # release previous video writer
                #             if vid_cap:  # video
                #                 fps = vid_cap.get(cv2.CAP_PROP_FPS)
                #                 w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
                #                 h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
                #             else:  # stream
                #                 fps, w, h = 30, im0.shape[1], im0.shape[0]
                #                 save_path += '.mp4'
                #             vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
                #         vid_writer[i].write(im0)

            # if cv2.waitKey(25) & self.stopEvent.is_set() == True:
            if cv2.waitKey(25) & self.in_detect == False:
                # self.stopEvent.clear()
                self.reset()
                break

        # Print results
        t = tuple(x / seen * 1E3 for x in dt)  # speeds per image
        LOGGER.info(
            f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}' % t)
        # if save_txt or save_img:
        #     s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
        #     LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
        # if update:
        #     strip_optimizer(weights)  # update model (to fix SourceChangeWarning)

    # 打开串口
    def port_open(self):
        try:
            # 配置串口参数
            # 打开 YAML 文件
            with open('config.yaml', 'r') as file:
                # 使用 PyYAML 加载 YAML 数据
                data = yaml.safe_load(file)
            self.ser.port = data['serial']  # 串口号
            self.ser.baudrate = 9600  # 波特率
            self.ser.bytesize = serial.EIGHTBITS  # 数据位，这里使用8位数据位
            self.ser.stopbits = serial.STOPBITS_ONE  # 停止位，这里使用1位停止位
            self.ser.parity = serial.PARITY_NONE  # 校验位，这里不使用校验位
            # 打开串口连接
            self.ser.open()
            if self.ser.is_open:
                print("串口已打开")
                return True
            else:
                do = input('串口打开失败，请输入真确的串口名，输入q否则退出')
                if do == 'q':
                    return False
                else:
                    data['serial'] = do
                    with open('config.yaml', 'w') as file:
                        yaml.dump(data, file)
                    self.port_open()
                return False
        except Exception as e:
            print(f"串口打开失败: {str(e)}")
            do = input('串口打开失败，请输入正确的串口名，输入q否则直接跳过串口打开')
            if do == 'q':
                return False
            else:
                data['serial'] = do
                with open('config.yaml', 'w') as file:
                    yaml.dump(data, file)
                self.port_open()
            return None
        # 打开串口接收定时器，周期为2ms
        self.timer.start(2)


if __name__ == '__main__':
    app = QtWidgets.QApplication(sys.argv)
    myshow = MainWindows()
    myshow.show()
    sys.exit(app.exec_())
